# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : Lv Shenkai
# @Time    : 2022/1/19 21:55
# @File    : cluster_user.py

import pandas
from sklearn.cluster import KMeans


def user_cluster(user_data_path, k):
    """
        通过用户属性进行K-means聚类
        对user进行分类
    :param user_data_path: 文件地址
    :param k: 聚类数
    :return: 返回user及其对应类别号的字典
    """

    # 读取文件，用0进行填充空值
    user_df = pandas.read_csv(user_data_path).fillna(value=0)

    # 获取user的属性列名(除去'客户id')
    column_name = list(user_df)
    column_name.remove('user_id')
    # print(column_name)

    # 进行数据规一化
    user_df[column_name] = \
        (user_df[column_name] - user_df.min()[column_name]) / \
        (user_df.max()[column_name] - user_df.min()[column_name])
    # 全0的数据列进行规一化会出现NaN值，用0进行填充
    user_df = user_df.fillna(value=0)

    # 进行K-Means聚类
    clf_KMeans = KMeans(n_clusters=k)
    # 得到聚类结果
    cluster = clf_KMeans.fit_predict(user_df[column_name])
    # print(cluster)

    # 聚类结果插入DataFrame中
    user_df.insert(1, 'user_type', cluster)
    # print(user_df)

    # 返回user_id和对应的类别
    # return user_df.iloc[:, 0:2]

    # 转化为字典
    user_type_dic = {}
    for index in range(len(cluster)):
        user_type_dic[index] = cluster[index]

    return user_type_dic
